use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;

#[derive(Debug)]
enum LSOutput {
    File(String, usize),
    Dir(String),
}

impl From<&str> for LSOutput {
    fn from(value: &str) -> Self {
        match &value[..3] {
            "dir" => LSOutput::Dir(value[4..].to_string()),
            _ => {
                let mut iter = value.split_whitespace();
                let size = iter.next().unwrap();
                let name = iter.next().unwrap();
                LSOutput::File(
                    name.to_string(),
                    size.parse::<usize>().unwrap()
                )
            }
        }
    }
}

#[derive(Debug)]
enum Command {
    CD(String),
    LS(Vec<LSOutput>),
}

#[derive(Debug)]
enum InodeKind {
    File,
    Dir(HashMap<String, Rc<RefCell<Inode>>>)
}

#[derive(Debug)]
struct Inode {
    size: usize,
    parent: Option<Rc<RefCell<Inode>>>,
    kind: InodeKind,
}

impl Inode {
    fn sub_inodes_mut(&mut self) -> &mut HashMap<String, Rc<RefCell<Inode>>> {
        match &mut self.kind {
            InodeKind::File => panic!(""),
            InodeKind::Dir(sub_inodes) => sub_inodes,
        }
    }

    fn sub_inodes(&self) -> &HashMap<String, Rc<RefCell<Inode>>> {
        match &self.kind {
            InodeKind::File => panic!(""),
            InodeKind::Dir(sub_inodes) => sub_inodes,
        }
    }

    fn parent(&self) -> Rc<RefCell<Inode>> {
        self.parent.as_ref().unwrap().clone()
    }

    fn calculate_size(&mut self) -> usize {
        match &mut self.kind {
            InodeKind::File => self.size,
            InodeKind::Dir(sub_inodes) => {
                let sub_inodes_size = sub_inodes
                    .values()
                    .map(|inode| inode.borrow_mut().calculate_size())
                    .sum::<usize>();
                self.size = sub_inodes_size;
                return self.size;
            }
        }
    }

    fn count_dir_size_at_most(inode: &Rc<RefCell<Inode>>, target_size: usize, record: &mut Vec<Rc<RefCell<Inode>>>) {
        if let InodeKind::Dir(sub_inodes) = &inode.borrow().kind {
            sub_inodes
                .values()
                .for_each(|inode| Inode::count_dir_size_at_most(inode, target_size, record));
            if inode.borrow().size <= target_size {
                record.push(inode.clone());
            }
        }
    }

    fn all_dir(inode: &Rc<RefCell<Inode>>) -> Vec<Rc<RefCell<Inode>>> {
        let mut dirs = Vec::new();
        fn do_all_dir(inode: &Rc<RefCell<Inode>>, dirs: &mut Vec<Rc<RefCell<Inode>>>) {
            if let InodeKind::Dir(sub_inodes) = &inode.borrow().kind {
                sub_inodes
                    .values()
                    .for_each(|inode| do_all_dir(inode, dirs));
                dirs.push(inode.clone());
            }
        }
        do_all_dir(inode, &mut dirs);
        dirs
    }
}

impl Inode {
    fn dir(parent: Option<Rc<RefCell<Inode>>>) -> Rc<RefCell<Self>> {
        Rc::new(RefCell::new(
            Inode {
                size: 0,
                parent,
                kind: InodeKind::Dir(HashMap::new()),
            }
        ))
    }
    
    fn file(size: usize, parent: Option<Rc<RefCell<Inode>>>) -> Rc<RefCell<Self>> {
        Rc::new(RefCell::new(
            Inode {
                size: size,
                parent,
                kind: InodeKind::File,
            }
        ))
    }

    fn from_commands(commands: Vec<Command>) -> Rc<RefCell<Inode>> {
        let root = Inode::dir(None);
        let mut current = root.clone();

        for command in commands.into_iter().skip(1) {
            match command {
                Command::CD(name) => {
                    let next = if name == ".." {
                        current.borrow_mut().parent()
                    } else {
                        current.borrow_mut().sub_inodes().get(&name).unwrap().clone()
                    };
                    // current must be a dir
                    current = next;
                }
                Command::LS(ls_outputs) => {
                    for ls_output in ls_outputs {
                        match ls_output {
                            LSOutput::File(name, size) => {
                                current.borrow_mut().sub_inodes_mut().insert(
                                    name.clone(),
                                    Inode::file(size, Some(current.clone()))   
                                );
                            }
                            LSOutput::Dir(name) => {
                                current.borrow_mut().sub_inodes_mut().insert(
                                    name.clone(),
                                    Inode::dir(Some(current.clone()))
                                );
                            }
                        }
                    }
                }
            }
        }

        root.borrow_mut().calculate_size();

        root
    }
}

fn parse_input() -> Rc<RefCell<Inode>> {
    let input = include_str!("../input.txt");
    let mut commands = Vec::new();
    let lines: Vec<_> = input.lines().into_iter().collect();

    let mut line_index = 0;
    while line_index < lines.len() {
        let line = lines[line_index];
        line_index += 1;
        // cd or ls?
        if line.starts_with("$ cd") {
            commands.push(Command::CD(line[5..].to_string()));
        } else {
            let mut ls_outputs: Vec<LSOutput> = Vec::new();
            // Loop util next command
            loop {
                if line_index < lines.len() && !lines[line_index].starts_with('$') {
                    ls_outputs.push(lines[line_index].into());
                    line_index += 1;
                } else {
                    break;
                }
            }
            commands.push(Command::LS(ls_outputs));
        }
    }

    Inode::from_commands(commands)
}

fn part1() {
    let root = parse_input();
    let mut record = Vec::new();
    Inode::count_dir_size_at_most(&root, 100000,&mut record);
    let sum = record
        .into_iter()
        .map(|inode| inode.borrow().size)
        .sum::<usize>();

    println!("Part 1: {}", sum);
}

fn part2() {
    let root = parse_input();
    let mut dirs = Inode::all_dir(&root);
    dirs.sort_by_key(|inode| inode.borrow().size);
    let least_dir_size = 30000000 - (70000000 - root.borrow().size);

    for inode in dirs {
        if inode.borrow().size >= least_dir_size {
            println!("Part 2: {}", inode.borrow().size);
            break;
        }
    }
}

fn main() {
    part1();
    part2();
}
